import json

from python_depend.prompt2data_api import data_query
from python_depend.redis_pool import RedisPoolManager

# ETF总资产排名
etf_assets_id = "685d1188901b6905dcc979b8"
# 股票总市值排名
stock_market_cap_id = "685e37ba3d66406735d74743"
source = "ETF-homepage"
stock_field_key = "国际美股@stock code_hq_code"
etf_field_key = '国际北美etf@etf code_hq_code'
redis_key_stock = "ds:magical_waves:stock_market_cap"
redis_key_etf = "ds:magical_waves:etf_assets"

# ==============================
# redis初始化
# ==============================
RedisPoolManager.initialize()
redis_client = RedisPoolManager.get('default')


def set_stock_data_cache():
    print(f"开始写入股票排名缓存数据")
    # 获取ETF总资产排名数据
    params = {
        "prompt_id": stock_market_cap_id,
        "data_limit": -1,
        "response_type": "add_hq_code",
        "source": source
    }
    result = data_query(params)
    if result is None:
        return ""
    print(f"查询结果:{str(result)}")
    stock_datas = result.get('data', {}).get("query_result", {}).get('datas', [{}])
    write_redis(redis_key_stock, stock_datas, stock_field_key)


def write_redis(key, stock_datas, field):
    result_data = []
    for stock_data in stock_datas:
        try:
            code_hq_code = stock_data.get(field)
            if code_hq_code is not None:
                market = code_hq_code.split(":")[0]
                code = code_hq_code.split(":")[1]
                market_code = {
                    "market": market,
                    "code": code
                }
                result_data.append(market_code)
        except Exception as e:
            print(f"处理数据异常:{str(stock_data)}", e)
            continue
    print(f"查询到数据:{str(result_data)}")
    list_json = json.dumps(result_data)
    redis_client.set(key, list_json)
    print(f"数据写入redis成功,key is {key}")


def set_etf_data_cache():
    print(f"开始写入etf总资产缓存数据")
    # 获取ETF总资产排名数据
    params = {
        "prompt_id": etf_assets_id,
        "data_limit": -1,
        "response_type": "add_hq_code",
        "source": source
    }
    result = data_query(params)
    if result.get('data') is None:
        return ""
    print(f"查询结果:{str(result)}")
    stock_datas = result.get('data', {}).get("query_result", {}).get('datas', [{}])
    num = int(len(stock_datas) / 10)
    stock_list = stock_datas[:num]
    write_redis(redis_key_etf, stock_list, etf_field_key)


# ==============================
# 主执行函数
# ==============================
def execute_job():
    set_stock_data_cache()
    set_etf_data_cache()
    redis_client.close()


if __name__ == '__main__':
    execute_job()
